# 1. 依赖库导入
import numpy as np
import math
import pygad
import matplotlib.pyplot as plt
import pandas as pd
from openpyxl import Workbook

# 绘图参数配置（支持中文显示）
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False
plt.switch_backend('TkAgg')

# 2. 题目核心参数定义（严格源自A题.pdf）
### 2.1 目标参数（🔶3-6、🔶3-8）
target_fake = np.array([0, 0, 0])  # 假目标：原点
target_true_center = np.array([0, 200, 0])  # 真目标下底面圆心
target_true_radius = 7  # 真目标圆柱半径（m）
target_true_height = 10  # 真目标圆柱高度（m）

### 2.2 导弹M1参数（🔶3-6、🔶3-5）
missile_M1_init = np.array([20000, 0, 2000])  # M1初始位置
missile_speed = 300  # M1飞行速度（m/s）
# 精确计算M1到达假目标的时间（≈66.999s）
missile_dist_to_fake = np.linalg.norm(target_fake - missile_M1_init)
missile_time_to_fake = missile_dist_to_fake / missile_speed  # 66.999s（核心约束）

### 2.3 无人机FY1参数（🔶3-6、🔶3-7）
drone_FY1_init = np.array([17800, 0, 1800])  # FY1初始位置
drone_speed_min = 70  # FY1最小速度（m/s）
drone_speed_max = 140  # FY1最大速度（m/s）
drone_flight_height = drone_FY1_init[2]  # FY1等高度飞行（z=1800m）

### 2.4 烟幕干扰弹参数（🔶3-5、🔶3-8）
smoke_radius = 10  # 烟幕有效半径（中心10m内）
smoke_sink_speed = 3  # 烟幕匀速下沉速度（m/s）
smoke_valid_duration = 20  # 烟幕有效时长（起爆后20s）
g = 9.8  # 重力加速度（m/s²）
smoke_drop_interval_min = 1  # 多弹投放最小间隔（≥1s）
smoke_min_explode_z = 50  # 起爆点最低z坐标（避免地面以下）


# 3. 核心工具函数（修复后）
### 3.1 真目标外接长方体顶点生成（🔶3-8）
def get_true_target_bounding_box():
    x_min = target_true_center[0] - target_true_radius
    x_max = target_true_center[0] + target_true_radius
    y_min = target_true_center[1] - target_true_radius
    y_max = target_true_center[1] + target_true_radius
    z_min = target_true_center[2]
    z_max = target_true_center[2] + target_true_height
    return [
        np.array([x_min, y_min, z_min]), np.array([x_max, y_min, z_min]),
        np.array([x_max, y_max, z_min]), np.array([x_min, y_max, z_min]),
        np.array([x_min, y_min, z_max]), np.array([x_max, y_min, z_max]),
        np.array([x_max, y_max, z_max]), np.array([x_min, y_max, z_max])
    ]


target_true_vertices = get_true_target_bounding_box()


### 3.2 导弹M1实时位置计算（🔶3-6）
def calculate_missile_real_time_pos(t):
    if t >= missile_time_to_fake - 1e-6:  # 到达后停止
        return target_fake
    flight_dir = target_fake - missile_M1_init
    flight_dir_unit = flight_dir / np.linalg.norm(flight_dir)
    displacement = missile_speed * t * flight_dir_unit
    return missile_M1_init + displacement


### 3.3 烟幕弹投放点/起爆点/实时中心计算（修复z坐标）
def calculate_smoke_drop_point(drone_speed, drone_heading, drop_delay):
    x = drone_FY1_init[0] + drone_speed * math.cos(drone_heading) * drop_delay
    y = drone_FY1_init[1] + drone_speed * math.sin(drone_heading) * drop_delay
    z = drone_flight_height
    return np.round(np.array([x, y, z]), 2)


def calculate_smoke_explosion_point(drone_speed, drone_heading, drop_delay, detonate_delay):
    drop_point = calculate_smoke_drop_point(drone_speed, drone_heading, drop_delay)
    # 修正：限制下落距离，确保起爆点z≥smoke_min_explode_z
    max_fall_distance = drop_point[2] - smoke_min_explode_z
    if max_fall_distance < 0:
        detonate_delay = 0.1  # 强制最小起爆延迟，避免z过低
        fall_distance = 0.5 * g * (detonate_delay ** 2)
    else:
        # 计算最大允许的起爆延迟（避免下落过多）
        max_allowed_det = math.sqrt(2 * max_fall_distance / g)
        detonate_delay = min(detonate_delay, max_allowed_det)
        fall_distance = 0.5 * g * (detonate_delay ** 2)

    x = drop_point[0] + drone_speed * math.cos(drone_heading) * detonate_delay
    y = drop_point[1] + drone_speed * math.sin(drone_heading) * detonate_delay
    z = max(drop_point[2] - fall_distance, smoke_min_explode_z)  # 确保z≥50
    return np.round(np.array([x, y, z]), 2)


def calculate_smoke_real_time_center(drone_speed, drone_heading, drop_delay, detonate_delay, t):
    detonate_time = drop_delay + detonate_delay
    if t < detonate_time - 1e-6:
        return np.array([0, 0, 0])
    explosion_point = calculate_smoke_explosion_point(drone_speed, drone_heading, drop_delay, detonate_delay)
    sink_time = t - detonate_time
    sink_distance = smoke_sink_speed * sink_time
    real_z = max(explosion_point[2] - sink_distance, 0.0)  # 下沉后不低于地面
    return np.array([explosion_point[0], explosion_point[1], real_z])


### 3.4 遮蔽有效性判断（🔶3-5、🔶3-8）
def is_smoke_effective(drone_speed, drone_heading, drop_delay, detonate_delay, t):
    missile_pos = calculate_missile_real_time_pos(t)
    smoke_center = calculate_smoke_real_time_center(drone_speed, drone_heading, drop_delay, detonate_delay, t)
    if np.all(smoke_center == np.array([0, 0, 0])):
        return False

    # 检查导弹-真目标顶点线段是否与烟幕球相交
    for vertex in target_true_vertices:
        vec_AC = smoke_center - missile_pos
        vec_AB = vertex - missile_pos
        ab_sq_len = np.dot(vec_AB, vec_AB)
        if ab_sq_len < 1e-12:
            continue

        k = np.dot(vec_AC, vec_AB) / ab_sq_len
        if k <= 0:
            dist = np.linalg.norm(vec_AC)
        elif k >= 1:
            dist = np.linalg.norm(smoke_center - vertex)
        else:
            cross_prod = np.cross(vec_AC, vec_AB)
            dist = np.linalg.norm(cross_prod) / math.sqrt(ab_sq_len)

        if dist <= smoke_radius + 1e-6:
            return True
    return False


### 3.5 多枚烟幕弹有效时间段合并（🔶3-8）
def get_single_smoke_effective_intervals(drone_speed, drone_heading, drop_delay, detonate_delay, time_step=0.05):
    detonate_time = drop_delay + detonate_delay
    effective_start = detonate_time
    effective_end = min(detonate_time + smoke_valid_duration, missile_time_to_fake)
    if effective_start >= effective_end + 1e-6:
        return []

    intervals = []
    current_start = None
    t = effective_start
    while t <= effective_end + 1e-6:
        if is_smoke_effective(drone_speed, drone_heading, drop_delay, detonate_delay, t):
            if current_start is None:
                current_start = t
        else:
            if current_start is not None:
                intervals.append((round(current_start, 2), round(t, 2)))
                current_start = None
        t += time_step

    if current_start is not None:
        intervals.append((round(current_start, 2), round(effective_end, 2)))
    return [(s, e) for s, e in intervals if (e - s) > 0.1]


def merge_smoke_intervals(all_intervals):
    flat_intervals = []
    for intervals in all_intervals:
        flat_intervals.extend(intervals)
    if not flat_intervals:
        return []
    sorted_intervals = sorted(flat_intervals, key=lambda x: x[0])
    merged = [sorted_intervals[0]]
    for current in sorted_intervals[1:]:
        last = merged[-1]
        if current[0] <= last[1] + 0.1:
            merged[-1] = (last[0], max(last[1], current[1]))
        else:
            merged.append(current)
    return merged


def calculate_total_effective_duration(params, time_step=0.05):
    v, theta, t1, t2, t3, d1, d2, d3 = params
    # 确保投放延迟有序（避免间隔为负）
    t1, t2, t3 = sorted([t1, t2, t3])
    interval1 = get_single_smoke_effective_intervals(v, theta, t1, d1, time_step)
    interval2 = get_single_smoke_effective_intervals(v, theta, t2, d2, time_step)
    interval3 = get_single_smoke_effective_intervals(v, theta, t3, d3, time_step)
    merged_intervals = merge_smoke_intervals([interval1, interval2, interval3])
    total_duration = sum([round(e - s, 2) for s, e in merged_intervals])
    return total_duration, merged_intervals, [interval1, interval2, interval3]


# 4. 遗传算法优化（修复约束检查）
### 4.1 适应度函数（严格约束）
def fitness_function(ga_instance, solution, solution_idx):
    v, theta, t1, t2, t3, d1, d2, d3 = solution

    # 约束1：速度范围（70≤v≤140）
    if not (drone_speed_min - 1e-6 <= v <= drone_speed_max + 1e-6):
        return 0.0

    # 约束2：投放延迟非负且有序（确保间隔≥1s）
    t_list = sorted([t1, t2, t3])
    if t_list[0] < -1e-6 or t_list[1] < -1e-6 or t_list[2] < -1e-6:
        return 0.0
    if (t_list[1] - t_list[0] < smoke_drop_interval_min - 1e-6 or
            t_list[2] - t_list[1] < smoke_drop_interval_min - 1e-6):
        return 0.0

    # 约束3：起爆延迟非负（≥0.1s）
    if d1 < 0.1 - 1e-6 or d2 < 0.1 - 1e-6 or d3 < 0.1 - 1e-6:
        return 0.0

    # 约束4：起爆时间≤导弹到达时间（≤66.999s）
    det_time1 = t_list[0] + d1
    det_time2 = t_list[1] + d2
    det_time3 = t_list[2] + d3
    if (det_time1 > missile_time_to_fake + 1e-6 or
            det_time2 > missile_time_to_fake + 1e-6 or
            det_time3 > missile_time_to_fake + 1e-6):
        return 0.0

    # 计算总有效时长
    total_duration, _, _ = calculate_total_effective_duration([v, theta, t_list[0], t_list[1], t_list[2], d1, d2, d3])
    return total_duration


### 4.2 迭代过程回调
generation_best_fitness = []


def on_generation_callback(ga_instance):
    current_best = np.max(ga_instance.last_generation_fitness)
    generation_best_fitness.append(current_best)
    print(f"迭代进度：{ga_instance.generations_completed:3d}代 | 当前最大遮蔽时长：{current_best:6.2f}s")


# 5. 主函数（修复后）
if __name__ == "__main__":
    # 5.1 遗传算法参数（优化搜索范围）
    gene_space = [
        {'low': drone_speed_min, 'high': drone_speed_max},  # 0: 速度v
        {'low': 0, 'high': 2 * math.pi},  # 1: 航向theta
        {'low': 0, 'high': 50},  # 2: t1（≤50s，避免超导弹时间）
        {'low': 0, 'high': 50},  # 3: t2
        {'low': 0, 'high': 50},  # 4: t3
        {'low': 0.1, 'high': 20},  # 5: d1（≤20s，避免下落过多）
        {'low': 0.1, 'high': 20},  # 6: d2
        {'low': 0.1, 'high': 20}  # 7: d3
    ]

    # 5.2 初始化遗传算法（增大种群，提高变异率）
    ga = pygad.GA(
        num_generations=200,  # 迭代次数
        num_parents_mating=25,  # 父代数量（增加多样性）
        fitness_func=fitness_function,
        sol_per_pop=150,  # 种群规模（从100→150）
        num_genes=8,
        gene_space=gene_space,
        parent_selection_type="sss",
        crossover_type="single_point",
        mutation_type="random",
        mutation_percent_genes=5,  # 变异率（从3%→5%）
        mutation_num_genes=1,
        parallel_processing=8,
        on_generation=on_generation_callback
    )

    # 5.3 启动优化
    print("=" * 60)
    print("无人机FY1投放3枚烟幕弹干扰M1策略优化（基于A题.pdf）")
    print(f"导弹到达假目标时间：{missile_time_to_fake:.2f}s | 烟幕最低起爆z：{smoke_min_explode_z}m")
    print("=" * 60)
    ga.run()
    print("=" * 60)
    print("优化完成！开始解析最优结果...")
    print("=" * 60)

    # 5.4 解析最优结果（确保投放延迟有序）
    best_solution, best_fitness, best_idx = ga.best_solution()
    v_opt = best_solution[0]
    theta_opt = best_solution[1]
    t1_raw, t2_raw, t3_raw = best_solution[2], best_solution[3], best_solution[4]
    d1_opt, d2_opt, d3_opt = best_solution[5], best_solution[6], best_solution[7]

    # 强制投放延迟有序（避免间隔问题）
    t_list = sorted([t1_raw, t2_raw, t3_raw])
    t1_opt, t2_opt, t3_opt = t_list[0], t_list[1], t_list[2]
    theta_opt_deg = (theta_opt * 180 / math.pi) % 360
    det1_opt = t1_opt + d1_opt
    det2_opt = t2_opt + d2_opt
    det3_opt = t3_opt + d3_opt

    # 计算关键位置（修复后z≥50m）
    drop1_opt = calculate_smoke_drop_point(v_opt, theta_opt, t1_opt)
    expl1_opt = calculate_smoke_explosion_point(v_opt, theta_opt, t1_opt, d1_opt)
    drop2_opt = calculate_smoke_drop_point(v_opt, theta_opt, t2_opt)
    expl2_opt = calculate_smoke_explosion_point(v_opt, theta_opt, t2_opt, d2_opt)
    drop3_opt = calculate_smoke_drop_point(v_opt, theta_opt, t3_opt)
    expl3_opt = calculate_smoke_explosion_point(v_opt, theta_opt, t3_opt, d3_opt)

    # 有效时间段计算
    total_duration_opt, merged_intervals_opt, single_intervals_opt = calculate_total_effective_duration(
        [v_opt, theta_opt, t1_opt, t2_opt, t3_opt, d1_opt, d2_opt, d3_opt]
    )
    fmt_interval = lambda intervals: ";".join([f"({s:.2f},{e:.2f})" for s, e in intervals]) if intervals else "无"
    interval1_str = fmt_interval(single_intervals_opt[0])
    interval2_str = fmt_interval(single_intervals_opt[1])
    interval3_str = fmt_interval(single_intervals_opt[2])
    merged_interval_str = fmt_interval(merged_intervals_opt)

    # 5.5 打印最优结果（验证约束满足）
    print("\n【最优投放策略总结】")
    print(f"1. 无人机FY1参数：")
    print(f"   - 飞行速度：{v_opt:.2f} m/s（70~140m/s，符合🔶3-7）")
    print(f"   - 飞行航向：{theta_opt_deg:.2f}°（x轴正向为0°）")
    print(f"   - 飞行高度：{drone_flight_height:.2f} m（等高度，符合🔶3-7）")
    print(f"\n2. 总有效遮蔽时长：{total_duration_opt:.2f} s（合并后，符合🔶3-8）")

    print(f"\n3. 第1枚烟幕弹参数：")
    print(f"   - 投放延迟：{t1_opt:.2f} s | 起爆延迟：{d1_opt:.2f} s")
    print(f"   - 起爆时间：{det1_opt:.2f} s（≤{missile_time_to_fake:.2f}s，符合🔶3-5）")
    print(f"   - 投放点：({drop1_opt[0]}, {drop1_opt[1]}, {drop1_opt[2]}) m")
    print(f"   - 起爆点：({expl1_opt[0]}, {expl1_opt[1]}, {expl1_opt[2]}) m（z≥50m）")
    print(f"   - 有效时间段：{interval1_str}（起爆后20s内，符合🔶3-5）")

    print(f"\n4. 第2枚烟幕弹参数：")
    print(f"   - 投放延迟：{t2_opt:.2f} s（与第1枚间隔{t2_opt - t1_opt:.2f}s ≥1s，符合🔶3-5）")
    print(f"   - 起爆延迟：{d2_opt:.2f} s | 起爆时间：{det2_opt:.2f} s")
    print(f"   - 投放点：({drop2_opt[0]}, {drop2_opt[1]}, {drop2_opt[2]}) m")
    print(f"   - 起爆点：({expl2_opt[0]}, {expl2_opt[1]}, {expl2_opt[2]}) m（z≥50m）")
    print(f"   - 有效时间段：{interval2_str}")

    print(f"\n5. 第3枚烟幕弹参数：")
    print(f"   - 投放延迟：{t3_opt:.2f} s（与第2枚间隔{t3_opt - t2_opt:.2f}s ≥1s，符合🔶3-5）")
    print(f"   - 起爆延迟：{d3_opt:.2f} s | 起爆时间：{det3_opt:.2f} s")
    print(f"   - 投放点：({drop3_opt[0]}, {drop3_opt[1]}, {drop3_opt[2]}) m")
    print(f"   - 起爆点：({expl3_opt[0]}, {expl3_opt[1]}, {expl3_opt[2]}) m（z≥50m）")
    print(f"   - 有效时间段：{interval3_str}")

    print(f"\n6. 合并后有效时间段：{merged_interval_str}（遮蔽可不连续，符合🔶3-8）")

    # 5.6 结果保存到result1.xlsx（🔶3-11要求）
    result_data = {
        "无人机编号": ["FY1", "FY1", "FY1"],
        "烟幕弹序号": [1, 2, 3],
        "飞行速度(m/s)": [round(v_opt, 2)] * 3,
        "飞行航向(°)": [round(theta_opt_deg, 2)] * 3,
        "投放延迟(s)": [round(t1_opt, 2), round(t2_opt, 2), round(t3_opt, 2)],
        "起爆延迟(s)": [round(d1_opt, 2), round(d2_opt, 2), round(d3_opt, 2)],
        "起爆时间(s)": [round(det1_opt, 2), round(det2_opt, 2), round(det3_opt, 2)],
        "投放点X(m)": [drop1_opt[0], drop2_opt[0], drop3_opt[0]],
        "投放点Y(m)": [drop1_opt[1], drop2_opt[1], drop3_opt[1]],
        "投放点Z(m)": [drop1_opt[2], drop2_opt[2], drop3_opt[2]],
        "起爆点X(m)": [expl1_opt[0], expl2_opt[0], expl3_opt[0]],
        "起爆点Y(m)": [expl1_opt[1], expl2_opt[1], expl3_opt[1]],
        "起爆点Z(m)": [expl1_opt[2], expl2_opt[2], expl3_opt[2]],
        "单弹有效时间段": [interval1_str, interval2_str, interval3_str],
        "单弹有效时长(s)": [
            round(sum(e - s for s, e in single_intervals_opt[0]), 2) if single_intervals_opt[0] else 0.0,
            round(sum(e - s for s, e in single_intervals_opt[1]), 2) if single_intervals_opt[1] else 0.0,
            round(sum(e - s for s, e in single_intervals_opt[2]), 2) if single_intervals_opt[2] else 0.0
        ],
        "总有效遮蔽时长(s)": [round(total_duration_opt, 2)] * 3,
        "合并后有效时间段": [merged_interval_str] * 3
    }

    df_result = pd.DataFrame(result_data)
    df_result.to_excel("result1.xlsx", index=False, engine="openpyxl")
    print(f"\n【结果保存】")
    print(f"最优策略已保存至 'result1.xlsx'（符合🔶3-11要求）")

    # 5.7 绘制迭代曲线
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, len(generation_best_fitness) + 1), generation_best_fitness,
             linewidth=2.5, marker='o', markersize=4, color='#2E86AB')
    plt.xlabel("遗传算法迭代次数", fontsize=12, fontweight='bold')
    plt.ylabel("最大有效遮蔽时长（s）", fontsize=12, fontweight='bold')
    plt.title("FY1投放3枚烟幕弹：优化迭代过程", fontsize=14, fontweight='bold')
    plt.grid(True, alpha=0.3, linestyle='--')
    plt.tight_layout()
    plt.savefig("FY1_3smoke_optimization_curve.png", dpi=300, bbox_inches='tight')
    plt.show()
    print(f"\n【可视化】")
    print(f"迭代过程曲线已保存至 'FY1_3smoke_optimization_curve.png'")
    print("=" * 60)